package com.example.streamer.util;

import javax.imageio.ImageIO;
import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.File;
import com.sun.jna.platform.win32.WinDef.HWND;
import com.sun.jna.platform.win32.WinDef.RECT;
import com.sun.jna.platform.win32.User32;
import com.sun.jna.platform.win32.GDI32;
import com.sun.jna.platform.win32.WinDef.HDC;
import com.sun.jna.platform.win32.WinDef.HBITMAP;
import com.sun.jna.platform.win32.WinGDI;
import com.sun.jna.Memory;
import com.sun.jna.platform.win32.WinNT.HANDLE;
import com.sun.jna.platform.win32.WinGDI.BITMAPINFO;
import com.sun.jna.Pointer;
import java.awt.Graphics2D;
import java.awt.Polygon;
import java.awt.MouseInfo;
import java.awt.Point;
import java.io.InputStream;

public class ScreenCapture {
    private static final Rectangle SCREEN_RECT = new Rectangle(Toolkit.getDefaultToolkit().getScreenSize());
    private static final Robot ROBOT;
    private static BufferedImage CURSOR_IMG = null;
    private static final int CURSOR_SIZE = 24;
    static {
        try {
            ROBOT = new Robot();
            InputStream is = ScreenCapture.class.getClassLoader().getResourceAsStream("cursor.png");
            if (is != null) {
                try {
                    CURSOR_IMG = ImageIO.read(is);
                } catch (Exception e) {
                    CURSOR_IMG = null;
                }
            }
        } catch (AWTException e) {
            throw new RuntimeException(e);
        }
    }

    public static byte[] captureScreenAsJpeg(float quality) {
        return captureScreenRectAsJpeg(SCREEN_RECT, quality);
    }

    public static byte[] captureScreenRectAsJpeg(Rectangle rect, float quality) {
        System.out.println("进入 ScreenCapture.captureScreenRectAsJpeg 方法，区域：" + rect);
        try {
            BufferedImage image = ROBOT.createScreenCapture(rect);
            // 叠加鼠标指针
            Point mouse = MouseInfo.getPointerInfo().getLocation();
            int mx = mouse.x - rect.x;
            int my = mouse.y - rect.y;
            if (mx >= 0 && my >= 0 && mx < rect.width && my < rect.height && CURSOR_IMG != null) {
                Graphics2D g = image.createGraphics();
                g.drawImage(CURSOR_IMG.getScaledInstance(CURSOR_SIZE, CURSOR_SIZE, java.awt.Image.SCALE_SMOOTH), mx, my, null);
                g.dispose();
            }
            String path = "G:/test.jpg";
            boolean ok = ImageIO.write(image, "jpg", new File(path));
            System.out.println("ImageIO.write 返回值: " + ok + "，保存路径: " + path);
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            ImageIO.write(image, "jpg", baos);
            return baos.toByteArray();
        } catch (Exception e) {
            System.out.println("抓屏异常！");
            e.printStackTrace();
            return null;
        }
    }

    public static byte[] captureWindowAsJpeg(HWND hwnd, float quality) {
        System.out.println("进入 ScreenCapture.captureWindowAsJpeg 方法，hwnd: " + hwnd);
        try {
            RECT rect = new RECT();
            User32.INSTANCE.GetWindowRect(hwnd, rect);
            int width = rect.right - rect.left;
            int height = rect.bottom - rect.top;
            HDC windowDC = User32.INSTANCE.GetDC(hwnd);
            HDC memDC = GDI32.INSTANCE.CreateCompatibleDC(windowDC);
            HBITMAP outputBitmap = GDI32.INSTANCE.CreateCompatibleBitmap(windowDC, width, height);
            HANDLE oldBitmap = GDI32.INSTANCE.SelectObject(memDC, outputBitmap);
            GDI32.INSTANCE.BitBlt(memDC, 0, 0, width, height, windowDC, 0, 0, 0x00CC0020);

            // 获取像素数据
            BITMAPINFO bmi = new BITMAPINFO();
            bmi.bmiHeader.biWidth = width;
            bmi.bmiHeader.biHeight = -height; // top-down
            bmi.bmiHeader.biPlanes = 1;
            bmi.bmiHeader.biBitCount = 32;
            bmi.bmiHeader.biCompression = WinGDI.BI_RGB;
            bmi.bmiHeader.biSize = bmi.bmiHeader.size();
            Memory buffer = new Memory(width * height * 4L);
            GDI32.INSTANCE.GetDIBits(memDC, outputBitmap, 0, height, buffer, bmi, WinGDI.DIB_RGB_COLORS);

            BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_RGB);
            int[] pixels = new int[width * height];
            buffer.read(0, pixels, 0, pixels.length);
            image.setRGB(0, 0, width, height, pixels, 0, width);

            // 叠加鼠标指针
            Point mouse = MouseInfo.getPointerInfo().getLocation();
            int mx = mouse.x - rect.left;
            int my = mouse.y - rect.top;
            if (mx >= 0 && my >= 0 && mx < width && my < height && CURSOR_IMG != null) {
                Graphics2D g = image.createGraphics();
                g.drawImage(CURSOR_IMG.getScaledInstance(CURSOR_SIZE, CURSOR_SIZE, java.awt.Image.SCALE_SMOOTH), mx, my, null);
                g.dispose();
            }

            GDI32.INSTANCE.SelectObject(memDC, oldBitmap);
            GDI32.INSTANCE.DeleteObject(outputBitmap);
            GDI32.INSTANCE.DeleteDC(memDC);
            User32.INSTANCE.ReleaseDC(hwnd, windowDC);

            ImageIO.write(image, "jpg", new File("G:/test.jpg"));
            ByteArrayOutputStream baos = new ByteArrayOutputStream();
            ImageIO.write(image, "jpg", baos);
            return baos.toByteArray();
        } catch (Exception e) {
            System.out.println("窗口抓取异常！");
            e.printStackTrace();
            return null;
        }
    }
} 